# light/core/logger.py
from abc import ABC

import torch # 导入torch，用于处理torch.Tensor类型
from pytorch_lightning.loggers.logger import Logger, rank_zero_experiment
from pytorch_lightning.utilities import rank_zero_only


class ConsoleLogger(Logger, ABC):
    def __init__(self, monitor='val_loss'):
        super().__init__()
        self.monitor = monitor
        self.value = None # 用于跟踪最佳值

    @property
    def name(self):
        return "ConsoleLogger" # 更改名称以更明确

    @property
    def version(self):
        return "0.1"

    @property
    @rank_zero_experiment
    def experiment(self):
        return None

    @rank_zero_only
    def log_hyperparams(self, params) -> None:
        print(f'\nlog_hyperparams: {params}')

    @rank_zero_only
    def log_metrics(self, metrics, step) -> None:
        # 创建一个副本以避免修改原始metrics字典
        current_metrics = metrics.copy()

        # 安全地获取epoch，如果不存在则默认为'N/A'
        epoch = current_metrics.pop('epoch', 'N/A') 

        formatted_metrics = []
        for k, v in current_metrics.items():
            # 处理torch.Tensor值：转换为标量Python数字
            if isinstance(v, torch.Tensor):
                if v.numel() == 1:
                    v = v.item()
                else:
                    # 如果是非标量张量，跳过或以其他方式处理
                    continue 
            
            # 根据类型格式化
            if isinstance(v, float):
                formatted_metrics.append(f'{k}: {v:.3f}')
            else: # 对于int, str等
                formatted_metrics.append(f'{k}: {v}')
        
        # 打印所有可用指标
        metric_text = ', '.join(formatted_metrics)
        print(f'  Metrics: epoch {epoch}, step {step}\n{{{metric_text}}}')

        # 安全地检查并更新监控值
        if self.monitor is not None and self.monitor in current_metrics: # <--- 关键检查在这里
            value = current_metrics[self.monitor]
            if isinstance(value, torch.Tensor) and value.numel() == 1:
                value = value.item() # 将张量转换为Python标量
            elif isinstance(value, torch.Tensor):
                # 如果是非标量张量，我们无法直接监控它，跳过
                return 

            # 假设'min'模式用于监控（例如val_loss）
            if (self.value is None) or (value < self.value): 
                self.value = value
                print(f' [Epoch {epoch}], [best {self.monitor}: {self.value:.3f}]')

    @rank_zero_only
    def save(self) -> None:
        super().save()

    @rank_zero_only
    def finalize(self, status: str) -> None:
        self.save()